# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import logging

from coin.strategy.mm.subscription import SubscriptionRequest
from coin.feed.fastfeed import feed_pb2
from coin.feed.fastfeed import proto_feed_gen

try:
  from coin.feed.fastfeed.native_wrapper import ProtoFeedBuilder
except (ImportError, ModuleNotFoundError):
  logging.warning('Unable to import native module.')


def create_proto_feed_handler(sub_req, feed_handler, *, no_parse=False, **kwargs):
  if sub_req == SubscriptionRequest('Futures', 'Bitmex'):
    return BitmexProtoFeedHandler(sub_req, feed_handler, no_parse, **kwargs)
  elif (sub_req == SubscriptionRequest('Futures', 'Okex', 'v3')
        or sub_req == SubscriptionRequest('Futures', 'Okex', 'v3_realtime')
        or sub_req == SubscriptionRequest('Futures', 'Okex', 'swap')):
    return OkexFuturesV3ProtoFeedHandler(sub_req, feed_handler, no_parse, **kwargs)
  else:
    return BookTradeProtoFeedHandler(sub_req, feed_handler, no_parse, **kwargs)


class ProtoFeedHandlerBase:
  def __init__(self, sub_req, feed_handler, no_parse=False):
    self._sub_req = sub_req
    self._feed_handler = feed_handler
    self._proto_feed_builder = ProtoFeedBuilder(sub_req=sub_req) if not no_parse else None
    self._proto_feed_callbacks = []
    self._published_metadata = set()
    if not no_parse:
      self.subscribe_proto_feed(self._forward_proto_feed_to_feed_builder)

  @property
  def sub_req(self):
    return self._sub_req

  @property
  def cc_builder(self):
    return self._proto_feed_builder.cc_builder

  def subscribe_proto_feed(self, callback):
    self._proto_feed_callbacks.append(callback)

  def publish_proto_feed(self, product, msgs):
    for callback in self._proto_feed_callbacks:
      callback(product, msgs)

  def _forward_proto_feed_to_feed_builder(self, product, msgs):
    for msg in msgs:
      self._proto_feed_builder.update(msg.SerializeToString())

  def _publish_metadata(self, timestamp, product):
    feed_arb_param_msg = proto_feed_gen.generate_feed_arb_param(timestamp, self._sub_req, product)
    book_builder_param_msg = proto_feed_gen.generate_book_builder_param(timestamp, product)
    self.publish_proto_feed(product, [feed_arb_param_msg, book_builder_param_msg])
    self._published_metadata.add(product)

  def _on_proto_feed(self, product, msgs):
    if product not in self._published_metadata:
      self._publish_metadata(msgs[0].timestamp, product)
    self.publish_proto_feed(product, msgs)

  # User methods
  def subscribe_event(self, event_type, keys, callback):
    if event_type == 'book':
      self.subscribe_book(keys, callback)
    elif event_type == 'trade':
      self.subscribe_trade(keys, callback)
    else:
      raise NotImplementedError(event_type)

  def subscribe(self, feed_type, symbol, callback, custom_feed_type=None):
    raise NotImplementedError()

  def subscribe_book(self, symbols, callback):
    raise NotImplementedError()

  def subscribe_trade(self, symbols, callback):
    raise NotImplementedError()

  # Provides transparency to the underlying feed handler. Allows users to call other 'subscribe_*'
  # methods.
  def __getattr__(self, name):
    if name.startswith('subscribe_'):
      logging.warning('%s is not supported by fastfeed.')
      return getattr(self._feed_handler, name)
    raise AttributeError(name)


class BookTradeProtoFeedHandler(ProtoFeedHandlerBase):
  def __init__(self, sub_req, feed_handler, no_parse=False):
    super().__init__(sub_req, feed_handler, no_parse)
    self._book_gen = None
    self._trade_gen = None

  def subscribe(self, feed_type, symbols, callback, custom_feed_type=None):
    if feed_type == feed_pb2.FEED_TYPE_BOOK:
      self.subscribe_book(symbols, callback)
    elif feed_type == feed_pb2.FEED_TYPE_TRADE:
      self.subscribe_trade(symbols, callback)
    else:
      raise ValueError('%s: not supported' % str(feed_type))

  def subscribe_book(self, symbols, callback):
    if self._book_gen is None:
      self._book_gen = proto_feed_gen.BookGenerator(feed_pb2.FEED_TYPE_BOOK)
      self._book_gen.subscribe_proto_feed(self._on_proto_feed)
    self._book_gen.add_symbols(self._feed_handler, symbols)

    if self._proto_feed_builder and callback is not None:
      self._proto_feed_builder.subscribe_book(symbols, callback)

  def subscribe_trade(self, symbols, callback):
    if self._trade_gen is None:
      self._trade_gen = proto_feed_gen.TradeGenerator(feed_pb2.FEED_TYPE_TRADE)
      self._trade_gen.subscribe_proto_feed(self._on_proto_feed)
    self._trade_gen.add_symbols(self._feed_handler, symbols)

    if self._proto_feed_builder and callback is not None:
      self._proto_feed_builder.subscribe_book(symbols, callback)


class BitmexProtoFeedHandler(BookTradeProtoFeedHandler):
  def __init__(self, sub_req, feed_handler, no_parse=False):
    super().__init__(sub_req, feed_handler, no_parse)
    self._inst_gen = None

  def subscribe_instrument(self, symbols, callback):
    if self._inst_gen is None:
      self._inst_gen = proto_feed_gen.BitmexInstrumentGenerator(feed_pb2.FEED_TYPE_CUSTOM,
                                                                'bitmex_instrument')
      self._inst_gen.subscribe_proto_feed(self._on_proto_feed)
    self._inst_gen.add_symbols(self._feed_handler, symbols)
    if callback:
      self._feed_handler.subscribe_instrument(symbols, callback)


class OkexFuturesV3ProtoFeedHandler(BookTradeProtoFeedHandler):
  def __init__(self, sub_req, feed_handler, no_parse=False):
    super().__init__(sub_req, feed_handler, no_parse)
    self._inst_gen = None

  def subscribe_instrument(self, symbols, callback):
    if self._inst_gen is None:
      self._inst_gen = proto_feed_gen.OkexFuturesV3InstrumentGenerator(
          feed_pb2.FEED_TYPE_CUSTOM, 'okex_instrument')
      self._inst_gen.subscribe_proto_feed(self._on_proto_feed)
    self._inst_gen.add_symbols(self._feed_handler, symbols)
    if callback:
      self._feed_handler.subscribe_instrument(symbols, callback)
